﻿using System.Collections.Generic; 
using Csu.Modsim.ModsimModel;
using ASquared.SymbolicMath;
using ASquared.ModelOptimization;

namespace iwpm.optimization
{
    public class LagrangianHydro : OptiFunctionBase, IConditionalGradientSolvableFunction
    {
        // Instance variables 
        private Model _model;
        private HydroCalculator _dispHydroFxn, _nondispHydroFxn;
        private double _hydroweight;
        private double[] _lambda, _mu; // Lagrangian multipliers
        private double[] _energyTargets;
        private double[] _energyDeficit;

        // Public fields 
        public bool UseOptimalityConditionDecomp = false;

        // Properties
        /// <summary>Gets and sets the energy deficit supplied by the optimal power flow model (used in the objective function only when UseOptimalityConditionDecomp is true).</summary>
        /// <remarks>Must be the same length as Lambda.</remarks>
        public double[] EnergyDeficit
        {
            get
            {
                return _energyDeficit;
            }
            set
            {
                _energyDeficit = value;
            }
        }
        /// <summary>Gets and sets the current estimate of the final energy targets for the dispatchable reservoirs (used as a constraint in the objective function only when UseOptimalityConditionDecomp is true).</summary>
        /// <remarks>Must be the same length as Mu.</remarks>
        public double[] EnergyTargets
        {
            get
            {
                return _energyTargets;
            }
            set
            {
                _energyTargets = value;
            }
        }
        /// <summary>Gets and sets the weight given to the water system optimization.</summary>
        public double HydroWeight
        {
            get
            {
                return _hydroweight;
            }
            set
            {
                _hydroweight = value;
            }
        }
        /// <summary>Gets and sets the power balance lagrangian multipliers including nondispatchable reservoirs.</summary>
        public double[] Lambda
        {
            get
            {
                return _lambda;
            }
            set
            {
                _lambda = value;
            }
        }
        /// <summary>Gets and sets the hydropower balance lagrangian multipliers for dispatchable reservoirs.</summary>
        public double[] Mu
        {
            get
            {
                return _mu;
            }
            set
            {
                _mu = value;
            }
        }
        /// <summary>Gets and sets the objective function of this instance.</summary>
        public override Symbol fxn
        {
            get
            {
                if (base.SymbolVariables.Length == 0)
                    return base.fxn;

                // Get the hydropower portions of the objective function
                Symbol nondispSum = this.GetNonDispatchablePortion();
                Symbol dispSum = this.GetDispatchablePortion();

                // Get the MODSIM objective function without the hydropower variables 
                List<string> hydroVars = nondispSum.Variables(); 
                hydroVars.AddRange(dispSum.Variables()); 
                Symbol water = this.GetWaterPortion(hydroVars.ToArray());

                // Return the whole objective function.
                if (nondispSum != null && dispSum != null)
                    return water + nondispSum + dispSum;
                else if (nondispSum != null)
                    return water + nondispSum;
                else if (dispSum != null)
                    return water + dispSum;
                else
                    return water;
            }
        }

        // Constructor
        /// <summary>The default hydropower objective function coming from the overarching Lagrangian function.</summary>
        /// <param name="_model">The MODSIM _model.</param>
        public LagrangianHydro(Model model, LagrangianHydroInputs inputs)
            : this(model, inputs.dispatchableUnits, inputs.nondispatchableUnits, inputs.weight_hydro, inputs.initialLambdas, inputs.initialMus)
        {
        }
        /// <summary>The default hydropower objective function coming from the overarching Lagrangian function.</summary>
        /// <param name="_model">The MODSIM _model.</param>
        /// <param name="weight_hydro">The weight associated with the water system simulation.</param>
        /// <param name="dispatchableUnits">A string of hydropower unit names separated by '|' specifying the dispatachable units within the system.</param>
        /// <param name="nondispatchableUnits">A string of hydropower unit names separated by '|' specifying the nondispatachable units within the system.</param>
        /// <param name="initialLambdas">The initial Lagrangian multipliers that are multiplied against the power balance coupling constraint.</param>
        /// <param name="initialMus">The initial Lagrangian multipliers that are multiplied against the power difference between the hydropower produced at dispatchable units and power injected into the power system at those units.</param>
        public LagrangianHydro(Model model, string dispatchableUnits, string nondispatchableUnits, double weight_hydro, double[] initialLambdas, double[] initialMus)
            : base(model)
        {
            _model = model;
            _dispHydroFxn = new HydroCalculator(model, HydropowerUnit.GetHydroUnitsFromString(model, dispatchableUnits));
            _nondispHydroFxn = new HydroCalculator(model, HydropowerUnit.GetHydroUnitsFromString(model, nondispatchableUnits));
            _hydroweight = weight_hydro;
            _lambda = initialLambdas;
            _mu = initialMus;
        }

        // Methods
        private void SetupMask()
        {
            double[] mask = new double[_model.Links_All.Length];
            foreach (HydropowerUnit unit in _dispHydroFxn.Units)
                foreach (Link link in unit.FlowLinks)
                    mask[link.number - 1] = 1;
            foreach (HydropowerUnit unit in _nondispHydroFxn.Units)
                foreach (Link link in unit.FlowLinks)
                    mask[link.number - 1] = 1; 
            base.variableMask = mask;
        }
        /// <summary>Initializes the instance for optimization.</summary>
        /// <remarks>This is a good place to build the objective function.</remarks>
        public override void Initialize()
        {
            base.Initialize();
            _dispHydroFxn.Initialize();
            _nondispHydroFxn.Initialize();

            HydroPreparer hPrep = new HydroPreparer(_model, _dispHydroFxn.Units, DefineConstants.COST_LARGE);
            hPrep.Prepare();
            hPrep = new HydroPreparer(_model, _nondispHydroFxn.Units, DefineConstants.COST_LARGE);
            hPrep.Prepare();

            this.SetupMask(); 
        }
        private Symbol GetNonDispatchablePortion()
        {
            Symbol nondispSum = null;
            if (UseOptimalityConditionDecomp)
            {
                if (_lambda.Length != 0 && _energyDeficit.Length == _lambda.Length)
                {
                    nondispSum = _nondispHydroFxn.Energy(0);
                    for (int i = 1; i < _nondispHydroFxn.Units.Length; i++)
                        nondispSum += _nondispHydroFxn.Energy(i);
                    nondispSum /= _lambda.Length;

                    Symbol lambdaSum = _lambda[0] * (nondispSum + _energyDeficit[0]);
                    for (int i = 1; i < _lambda.Length; i++)
                        lambdaSum += _lambda[i] * (nondispSum + _energyDeficit[i]);

                    nondispSum = lambdaSum;
                }
            }
            else
            {
                if (_nondispHydroFxn.Units.Length != 0 && _lambda.Length != 0)
                {
                    Symbol lambdaSum = _lambda[0];
                    for (int i = 1; i < _lambda.Length; i++)
                        lambdaSum += _lambda[i];

                    nondispSum = _nondispHydroFxn.Energy(0);
                    for (int i = 1; i < _nondispHydroFxn.Units.Length; i++)
                        nondispSum += _nondispHydroFxn.Energy(i);
                    nondispSum = lambdaSum * nondispSum;
                }
            }
            return nondispSum;
        }
        private Symbol GetDispatchablePortion()
        {
            Symbol dispSum = null;
            if (UseOptimalityConditionDecomp)
            {
                if (_energyTargets.Length != 0 && _dispHydroFxn.Units.Length == _energyTargets.Length)
                {
                    dispSum = ((_dispHydroFxn.Energy(0) - _energyTargets[0]) ^ 2);
                    for (int i = 1; i < _dispHydroFxn.Units.Length; i++)
                        dispSum += ((_dispHydroFxn.Energy(i) - _energyTargets[i]) ^ 2);
                }
            }
            else
            {
                if (_mu.Length != 0 && _dispHydroFxn.Units.Length == _mu.Length)
                {
                    dispSum = _dispHydroFxn.Energy(0) * _mu[0];
                    for (int i = 1; i < _dispHydroFxn.Units.Length; i++)
                        dispSum += _dispHydroFxn.Energy(i) * _mu[i];
                }
            }
            return dispSum;
        }
        private Symbol GetWaterPortion(string[] excludeVariables)
        {
            return _hydroweight * base.GetModsimObjectiveFunction(excludeVariables);
        }
    }
}
